#!/usr/bin/env python3
# -*- coding: utf-8 -*-
# @Time    : 2024-05-21 10:36
# @Author  : nedncs
# @File    : RSADecrypt.py
# @Description : src/python_ldap_admin/utils
import base64

import Env
from Crypto import Random
from Crypto.PublicKey import RSA
from Crypto.Cipher import PKCS1_v1_5 as Cipher_pkcs1_v1_5

from core.common.logger import logger

private_key_file = "{}/go-ldap-admin-priv.pem".format(Env.DATA_DIR)
public_key_file = "{}/go-ldap-admin-pub.pem".format(Env.DATA_DIR)


def generate_rsa_key():
    # 伪随机数生成器
    random_generator = Random.new().read
    # rsa算法生成实例
    rsa = RSA.generate(1024, random_generator)
    private_file = private_key_file
    public_file = public_key_file
    # 私钥的生成
    private_pem = rsa.exportKey()
    with open(private_file, "wb") as f:
        f.write(private_pem)

    # 公钥的生成
    public_pem = rsa.publickey().exportKey()
    with open(public_file, "wb") as f:
        f.write(public_pem)


# rsa加密
@logger.catch
def rsa_encrypt(message):
    public_file = public_key_file
    with open(public_file, 'r') as f:
        public_key = f.read()

    rsa_key = RSA.importKey(public_key)
    cipher = Cipher_pkcs1_v1_5.new(rsa_key)  # 创建用于执行pkcs1_v1_5加密或解密的密码
    cipher_text = base64.b64encode(cipher.encrypt(message.encode("utf-8")))
    return cipher_text.decode("utf-8")


# rsa解密
@logger.catch
def rsa_decrypt(cipher_text):
    private_file = private_key_file
    with open(private_file, 'r') as f:
        private_key = f.read()

    decode_text = base64.b64decode(cipher_text)
    cipher = Cipher_pkcs1_v1_5.new(RSA.importKey(private_key))  # 创建用于执行pkcs1_v1_5加密或解密的密码
    decrypt_text = cipher.decrypt(decode_text, b"rsa")
    return decrypt_text.decode("utf-8")
